#
#   Richardson_RK4_SHO.py
#
#   Solves the Phase space dynamics of Simple Harmonic Oscillator
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from FunctionalBohm import SchrodingerNDT

# Problem setup

from HarmonicOscillator1D import Psi as wvfn
from HarmonicOscillator1D import V
from HarmonicOscillator1D import mass
from HarmonicOscillator1D import hbar
from HarmonicOscillator1D import omega

# Initial Conditions
from math import pi
timesteps = 200
T_initial = 0.5*pi/omega
T_final = T_initial + 2*(2*pi/omega) # one period

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

def csv(*args):
    if len(args) == 1:
        return str(args[0])+"\n"
    else:
        return str(args[0])+","+csv(*args[1:])

# Use RK4 with Richardson interpolation to propagate statistical ensembles:

S = SchrodingerNDT(V,wvfn,mass,dimensions=1,dx=0.000000001)

outfile = open("./RichardsonRK4.csv",'w')

for X_ic in segments(-10,10,25): # different starting places
    print(X_ic)
    traj = S.get_Trajectory(X_ic,T_initial,T_final,timesteps) # trajectory as a list.
    time = S.get_Times(T_initial,T_final,timesteps)
    #p = map(S.get_momentum(),time,traj)
    #Q = map(S.get_momentum_variance(),time,traj)
    for i in range(timesteps):
        outfile.write(csv(time[i],traj[i]))#,p[i][0])) #,Q[i]))

outfile.close()
print("Updated './RichardsonRK4.csv'")
